
#include <cuda.h>
#include <iostream>
#include <math_constants.h>
#include "math.hpp"
#include <opencv2/viz/vizcore.hpp>
#include "device_array.cuh"
#include "safe_call.hpp"
#include "raycast.h"
#include "raycast.cuh"
using namespace kfusion::Cuda;

#pragma pack(push, 1)

#pragma pack(pop)

namespace device
{
    struct TsdfVolume
    {
    public:
        typedef struct Vovel elem_type;
        elem_type *data;

        const int3 dims;
        const float3 voxel_size;
        const float trunc_dist;
        const int max_weight;

        TsdfVolume(int3 _dims, float3 _voxel_size, float _trunc_dist, int _max_weight)
            : dims(_dims), voxel_size(_voxel_size), trunc_dist(_trunc_dist), max_weight(_max_weight)
        {
            cudaMalloc(&data, sizeof(Vovel) * 512 * 512 * 512);
        };
        TsdfVolume(elem_type *data, int3 _dims, float3 _voxel_size, float _trunc_dist, int _max_weight)
            : data(data), dims(_dims), voxel_size(_voxel_size), trunc_dist(_trunc_dist), max_weight(_max_weight){};
        __host__ elem_type *ptr() const
        {
            return data;
        }
        __host__ void release() const
        {
            cudaFree(data);
        }
        __kf_device__ elem_type &operator()(int x, int y, int z);
        __kf_device__ const elem_type *operator()(int x, int y, int z) const;
        __kf_device__ elem_type *beg(int x, int y) const;
        __kf_device__ elem_type *zstep(elem_type *const ptr) const;
        __device__ float interpolate(float3 p_voxels) const
        {
            float3 cf = p_voxels;
            // TsdfVolume &volume = *(this);
            // rounding to negative infinity
            int3 g = make_int3(__float2int_rd(cf.x), __float2int_rd(cf.y), __float2int_rd(cf.z));

            if (g.x < 0 || g.x >= this->dims.x - 1 || g.y < 0 || g.y >= this->dims.y - 1 || g.z < 0 || g.z >= this->dims.z - 1)
                return CUDART_NAN_F;

            float a = cf.x - g.x;
            float b = cf.y - g.y;
            float c = cf.z - g.z;

            float tsdf = 0.f;
            tsdf += operator()(g.x + 0, g.y + 0, g.z + 0)->tsdf * (1 - a) * (1 - b) * (1 - c);
            tsdf += operator()(g.x + 0, g.y + 0, g.z + 1)->tsdf * (1 - a) * (1 - b) * c;
            tsdf += operator()(g.x + 0, g.y + 1, g.z + 0)->tsdf * (1 - a) * b * (1 - c);
            tsdf += operator()(g.x + 0, g.y + 1, g.z + 1)->tsdf * (1 - a) * b * c;
            tsdf += operator()(g.x + 1, g.y + 0, g.z + 0)->tsdf * a * (1 - b) * (1 - c);
            tsdf += operator()(g.x + 1, g.y + 0, g.z + 1)->tsdf * a * (1 - b) * c;
            tsdf += operator()(g.x + 1, g.y + 1, g.z + 0)->tsdf * a * b * (1 - c);
            tsdf += operator()(g.x + 1, g.y + 1, g.z + 1)->tsdf * a * b * c;
            return tsdf;
        }

    private:
        TsdfVolume &operator=(const TsdfVolume &);
    };

    __kf_device__ void intersect(float3 ray_org, float3 ray_dir, /*float3 box_min,*/ float3 box_max, float &tnear, float &tfar)
    {
        const float3 box_min = make_float3(0.f, 0.f, 0.f);

        // compute intersection of ray with all six bbox planes
        float3 invR = make_float3(1.f / ray_dir.x, 1.f / ray_dir.y, 1.f / ray_dir.z);
        float3 tbot = invR * (box_min - ray_org);
        float3 ttop = invR * (box_max - ray_org);

        // re-order intersections to find smallest and largest on each axis
        float3 tmin = make_float3(fminf(ttop.x, tbot.x), fminf(ttop.y, tbot.y), fminf(ttop.z, tbot.z));
        float3 tmax = make_float3(fmaxf(ttop.x, tbot.x), fmaxf(ttop.y, tbot.y), fmaxf(ttop.z, tbot.z));

        // find the largest tmin and the smallest tmax
        tnear = fmaxf(fmaxf(tmin.x, tmin.y), fmaxf(tmin.x, tmin.z));
        tfar = fminf(fminf(tmax.x, tmax.y), fminf(tmax.x, tmax.z));
    }

    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    /// TsdfVolume

    __kf_device__ TsdfVolume::elem_type &TsdfVolume::operator()(int x, int y, int z)
    {
        return data[x + y * dims.x + z * dims.y * dims.x];
    }

    __kf_device__ const TsdfVolume::elem_type *TsdfVolume::operator()(int x, int y, int z) const
    {
        return &data[x + y * dims.x + z * dims.y * dims.x];
    }

    __kf_device__ TsdfVolume::elem_type *TsdfVolume::beg(int x, int y) const
    {
        return &data[x + y * dims.x];
    }

    __kf_device__ TsdfVolume::elem_type *TsdfVolume::zstep(elem_type *const ptr) const
    {
        return ptr + dims.x * dims.y;
    }

    struct TsdfRaycaster
    {
        device::TsdfVolume volume;

        Aff3f aff;
        Mat3f Rinv;
        Camera reproj;
        Vec3f volume_size;
        float time_step;
        float3 gradient_delta;
        float3 voxel_size_inv;

        TsdfRaycaster(const TsdfVolume &volume, const Aff3f &aff, const Mat3f &Rinv, const Camera &_reproj);

        __kf_device__ float fetch_tsdf(const float3 &p) const
        {
            // rounding to nearest even
            int x = __float2int_rn(p.x * voxel_size_inv.x);
            int y = __float2int_rn(p.y * voxel_size_inv.y);
            int z = __float2int_rn(p.z * voxel_size_inv.z);
            TsdfVolume::elem_type vox = *volume(x, y, z);
            return vox.tsdf;
        }
        __kf_device__ void operator()(PtrStepSz<float3> points, PtrStepSz<float3> normals) const
        {
            int x = blockIdx.x * blockDim.x + threadIdx.x;
            int y = blockIdx.y * blockDim.y + threadIdx.y;

            if (x >= points.cols || y >= points.rows)
                return;
            const float qnan = CUDART_NAN_F;

            points(y, x) = normals(y, x) = make_float3(qnan, qnan, qnan);

            float3 ray_org = aff.t;
            float3 ray_dir = normalized(aff.R * reproj.cam2world(x, y, 1.f));
            // printf("%f,%f,%f,%f,%f,%f\n", ray_org.x, ray_org.y, ray_org.z, ray_dir.x, ray_dir.y, ray_dir.z);

            // We do subtract Vovel size to minimize checks after
            // Note: origin of volume coordinate is placeed
            // in the center of Vovel (0,0,0), not in the corener of the Vovel!
            float3 box_max = volume_size - volume.voxel_size;

            float tmin, tmax;
            intersect(ray_org, ray_dir, box_max, tmin, tmax);

            const float min_dist = 0.f;
            tmin = fmax(min_dist, tmin);
            if (tmin >= tmax)
                return;

            tmax -= time_step;
            float3 vstep = ray_dir * time_step;
            float3 next = ray_org + ray_dir * tmin;

            float tsdf_next = fetch_tsdf(next);
            for (float tcurr = tmin; tcurr < tmax; tcurr += time_step)
            {
                float tsdf_curr = tsdf_next;
                float3 curr = next;
                next += vstep;

                tsdf_next = fetch_tsdf(next);
                if (tsdf_curr < 0.f && tsdf_next > 0.f)
                    break;

                if (tsdf_curr > 0.f && tsdf_next < 0.f)
                {
                    float Ft = volume.interpolate(curr * voxel_size_inv);
                    float Ftdt = volume.interpolate(next * voxel_size_inv);

                    float Ts = tcurr - __fdividef(time_step * Ft, Ftdt - Ft);

                    float3 vertex = ray_org + ray_dir * Ts;
                    float3 normal = compute_normal(vertex);

                    if (!isnan(normal.x * normal.y * normal.z))
                    {
                        normal = Rinv * normal;
                        vertex = Rinv * (vertex - aff.t);

                        normals(y, x) = make_float3(normal.x, normal.y, normal.z);
                        points(y, x) = make_float3(vertex.x, vertex.y, vertex.z);
                    }
                    break;
                }
            } /* for (;;) */
        }

        __kf_device__ float3 compute_normal(const float3 &p) const
        {
            float3 n;

            float Fx1 = volume.interpolate(make_float3(p.x + gradient_delta.x, p.y, p.z) * voxel_size_inv);
            float Fx2 = volume.interpolate(make_float3(p.x - gradient_delta.x, p.y, p.z) * voxel_size_inv);
            n.x = __fdividef(Fx1 - Fx2, gradient_delta.x);

            float Fy1 = volume.interpolate(make_float3(p.x, p.y + gradient_delta.y, p.z) * voxel_size_inv);
            float Fy2 = volume.interpolate(make_float3(p.x, p.y - gradient_delta.y, p.z) * voxel_size_inv);
            n.y = __fdividef(Fy1 - Fy2, gradient_delta.y);

            float Fz1 = volume.interpolate(make_float3(p.x, p.y, p.z + gradient_delta.z) * voxel_size_inv);
            float Fz2 = volume.interpolate(make_float3(p.x, p.y, p.z - gradient_delta.z) * voxel_size_inv);
            n.z = __fdividef(Fz1 - Fz2, gradient_delta.z);

            return normalized(n);
        }
    };

    inline TsdfRaycaster::TsdfRaycaster(const TsdfVolume &_volume, const Aff3f &_aff, const Mat3f &_Rinv, const Camera &_reproj)
        : volume(_volume), aff(_aff), Rinv(_Rinv), reproj(_reproj) {}

    __global__ void raycast_kernel(const TsdfRaycaster raycaster, PtrStepSz<float3> points, PtrStepSz<float3> normals)
    {
        raycaster(points, normals);
    };

}
// }

__global__ void render_image_kernel(const PtrStepSz<float3> points, const PtrStepSz<float3> normals,
                                    const Intr reproj, const float3 light_pose, PtrStepSz<uchar3> dst)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= points.cols || y >= points.rows)
        return;

    float3 color;

    float3 p = device::tr<float3>(points(y, x));

    if (isnan(p.x))
    {
        const float3 bgr1 = make_float3(4.f / 255.f, 2.f / 255.f, 2.f / 255.f);
        const float3 bgr2 = make_float3(236.f / 255.f, 120.f / 255.f, 120.f / 255.f);

        float w = static_cast<float>(y) / dst.rows;
        color = bgr1 * (1 - w) + bgr2 * w;
    }
    else
    {
        float3 P = p;
        float3 N = device::tr<float3>(normals(y, x));

        const float Ka = 0.3f; // ambient coeff
        const float Kd = 0.5f; // diffuse coeff
        const float Ks = 0.2f; // specular coeff
        const float n = 20.f;  // specular power

        const float Ax = 1.f; // ambient color,  can be RGB
        const float Dx = 1.f; // diffuse color,  can be RGB
        const float Sx = 1.f; // specular color, can be RGB
        const float Lx = 1.f; // light color

        // Ix = Ax*Ka*Dx + Att*Lx [Kd*Dx*(N dot L) + Ks*Sx*(R dot V)^n]

        float3 L = normalized(light_pose - P);
        float3 V = normalized(make_float3(0.f, 0.f, 0.f) - P);
        float3 R = normalized(2 * N * dot(N, L) - L);

        float Ix = Ax * Ka * Dx + Lx * Kd * Dx * fmax(0.f, dot(N, L)) + Lx * Ks * Sx * __powf(fmax(0.f, dot(R, V)), n);
        color = make_float3(Ix, Ix, Ix);
    }

    uchar3 out;
    out.x = static_cast<unsigned char>(__saturatef(color.x) * 255.f);
    out.y = static_cast<unsigned char>(__saturatef(color.y) * 255.f);
    out.z = static_cast<unsigned char>(__saturatef(color.z) * 255.f);
    dst(y, x) = out;
}

template <typename D, typename S>
inline D device_cast(const S &source)
{
    return *reinterpret_cast<const D *>(source.val);
}
inline device::Aff3f device_castA(const cv::Affine3f &source)
{
    device::Aff3f aff;
    cv::Matx<float, 3, 3> R = source.rotation();
    cv::Vec3f t = source.translation();
    aff.R = device_cast<device::Mat3f>(R);
    aff.t = device_cast<device::Vec3f>(t);
    return aff;
}

__global__ void tangent_colors_kernel(PtrStepSz<Normal> normals, PtrStep<uchar3> colors)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= normals.cols || y >= normals.rows)
        return;

    Normal n = normals(y, x);

#if 0
            unsigned char r = static_cast<unsigned char>(__saturatef((-n.x + 1.f)/2.f) * 255.f);
            unsigned char g = static_cast<unsigned char>(__saturatef((-n.y + 1.f)/2.f) * 255.f);
            unsigned char b = static_cast<unsigned char>(__saturatef((-n.z + 1.f)/2.f) * 255.f);
#else
    unsigned char r = static_cast<unsigned char>((5.f - n.x * 3.5f) * 25.5f);
    unsigned char g = static_cast<unsigned char>((5.f - n.y * 2.5f) * 25.5f);
    unsigned char b = static_cast<unsigned char>((5.f - n.z * 3.5f) * 25.5f);
#endif
    colors(y, x) = make_uchar3(b, g, r);
}

#include "iostream"
#include "opencv2/opencv.hpp"

void raycast_add(cv::Mat &img, cv::Mat &nor, cv::Affine3f camera_pose, const void *pdata, float3 center, float4 cam)
{
    cv::Mat cloud, normalsm;

    Camera intr(cam);
    auto dims = make_int3(512, 512, 512);
    auto vsz = make_float3(VOXELSIZE, VOXELSIZE,VOXELSIZE);
    float trunc_dist_ = 0.03f;
    int max_weight_ = 256;
    float raycaster_step_factor = 0.15;

    float gradient_delta_factor = 0.15;
    kfusion::Cuda::DeviceArray2D<float3> points, norm;
    int col = img.cols, row = img.rows;
    points.create(row, col);
    norm.create(row, col);
    kfusion::Cuda::DeviceArray2D<uchar3> dst;
    dst.create(row, col);

    cv::Affine3f pose = cv::Affine3f().translate(cv::Vec3f(center.x, center.y, center.z));

    device::TsdfVolume volume((device::TsdfVolume::elem_type *)pdata, dims, vsz, trunc_dist_, max_weight_);

    cv::Affine3f cam2vol = pose.inv() * camera_pose;
    device::Aff3f aff = device_castA(cam2vol);
    device::Mat3f Rinv = device_cast<device::Mat3f>(cam2vol.rotation().t()); //旋转矩阵视为单位矩阵 逆矩阵等于转置矩阵

    device::TsdfRaycaster rc(volume, aff, Rinv, intr);

    rc.volume_size = volume.voxel_size * volume.dims;
    rc.time_step = volume.trunc_dist * raycaster_step_factor;
    rc.gradient_delta = volume.voxel_size * gradient_delta_factor;
    rc.voxel_size_inv = 1.f / volume.voxel_size;

    dim3 block(32, 32);
    dim3 grid(kfusion::Cuda::divUp(points.cols(), block.x), kfusion::Cuda::divUp(points.rows(), block.y));

    raycast_kernel<<<grid, block>>>(rc, (PtrStepSz<float3>)points, (PtrStepSz<float3>)norm);
    ck(cudaDeviceSynchronize());
    float3 lp = make_float3(0, 0, 0);
    render_image_kernel<<<grid, block>>>((PtrStepSz<float3>)points, (PtrStepSz<float3>)norm, intr, lp, dst);
    ck(cudaDeviceSynchronize());
    dst.download(img.ptr<void>(), img.step);

    tangent_colors_kernel<<<grid, block>>>(norm, dst);

    dst.download(nor.ptr<void>(), nor.step);

    dst.release();
    points.release();
    norm.release();
}
// #include "../"
#include <iostream>
#include <fstream>

int main()
{
    std::fstream file("abgrids.bin", std::ios::out | std::ios::binary); // | ios::app
    struct Voxel32 pboxs;
    size_t len = gpu_pbox_use.size();
    file.write(reinterpret_cast<char *>(&len), sizeof(size_t));

    for (int i = 0; i < gpu_pbox_use.size(); i++)
    {
        ck(cudaMemcpy((void *)&pboxs, (void *)(gpu_pbox_use[i]), sizeof(struct Voxel32), cudaMemcpyDeviceToHost));
        file.write(reinterpret_cast<char *>(&pboxs), sizeof(struct Voxel32));
    }
    std::cout << gpu_pbox_use.size() << std::endl;
    file.close();
}